"""
RTSP客户端测试工具
用于测试RTSP服务器连接和流质量
"""

import cv2
import time
import sys
import threading
from datetime import datetime

class RTSPClient:
    def __init__(self, rtsp_url):
        self.rtsp_url = rtsp_url
        self.cap = None
        self.frame_count = 0
        self.start_time = None
        self.running = False
        
    def test_connection(self):
        """测试RTSP连接"""
        print(f"测试RTSP连接: {self.rtsp_url}")
        
        try:
            # 尝试连接
            cap = cv2.VideoCapture(self.rtsp_url)
            if not cap.isOpened():
                print("❌ 无法连接到RTSP服务器")
                return False
                
            # 读取第一帧
            ret, frame = cap.read()
            if not ret:
                print("❌ 无法读取视频流")
                cap.release()
                return False
                
            # 获取流信息
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            
            print("✅ RTSP连接成功!")
            print(f"   分辨率: {width}x{height}")
            print(f"   帧率: {fps:.1f} FPS")
            print(f"   第一帧大小: {frame.shape}")
            
            cap.release()
            return True
            
        except Exception as e:
            print(f"❌ 连接测试失败: {e}")
            return False
            
    def start_preview(self):
        """启动实时预览"""
        print(f"启动RTSP预览: {self.rtsp_url}")
        print("按 'q' 键退出预览")
        
        self.cap = cv2.VideoCapture(self.rtsp_url)
        if not self.cap.isOpened():
            print("❌ 无法打开RTSP流")
            return
            
        self.running = True
        self.start_time = time.time()
        
        # 启动统计线程
        stats_thread = threading.Thread(target=self._print_stats, daemon=True)
        stats_thread.start()
        
        try:
            while self.running:
                ret, frame = self.cap.read()
                if not ret:
                    print("❌ 读取帧失败")
                    break
                    
                self.frame_count += 1
                
                # 在帧上添加信息
                self._draw_info(frame)
                
                # 显示帧
                cv2.imshow('RTSP Preview', frame)
                
                # 检查退出键
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
                    
        except KeyboardInterrupt:
            print("\n用户中断")
        finally:
            self.stop_preview()
            
    def stop_preview(self):
        """停止预览"""
        self.running = False
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()
        
        if self.start_time:
            duration = time.time() - self.start_time
            avg_fps = self.frame_count / duration if duration > 0 else 0
            print(f"\n统计信息:")
            print(f"   运行时间: {duration:.1f}秒")
            print(f"   总帧数: {self.frame_count}")
            print(f"   平均帧率: {avg_fps:.1f} FPS")
            
    def _draw_info(self, frame):
        """在帧上绘制信息"""
        if self.start_time:
            duration = time.time() - self.start_time
            fps = self.frame_count / duration if duration > 0 else 0
            
            # 绘制信息
            info_text = [
                f"Frame: {self.frame_count}",
                f"FPS: {fps:.1f}",
                f"Time: {datetime.now().strftime('%H:%M:%S')}",
                f"URL: {self.rtsp_url}"
            ]
            
            y_offset = 30
            for text in info_text:
                cv2.putText(frame, text, (10, y_offset), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                y_offset += 25
                
    def _print_stats(self):
        """打印统计信息"""
        last_count = 0
        while self.running:
            time.sleep(5)  # 每5秒打印一次
            if self.running:
                current_fps = (self.frame_count - last_count) / 5
                print(f"实时FPS: {current_fps:.1f}, 总帧数: {self.frame_count}")
                last_count = self.frame_count

def main():
    # 默认RTSP地址
    default_urls = [
        "rtsp://localhost:8554/stream",
        "rtsp://localhost:8554/camera", 
        "rtsp://localhost:8554/test"
    ]
    
    if len(sys.argv) > 1:
        rtsp_url = sys.argv[1]
    else:
        print("RTSP客户端测试工具")
        print("=" * 40)
        print("可用命令:")
        print("  python rtsp_client.py <rtsp_url>  - 测试指定RTSP地址")
        print("  python rtsp_client.py             - 测试默认地址")
        print()
        
        # 测试默认地址
        for url in default_urls:
            print(f"测试: {url}")
            client = RTSPClient(url)
            if client.test_connection():
                rtsp_url = url
                break
            print()
        else:
            print("❌ 所有默认RTSP地址都无法连接")
            print("请确保RTSP服务器正在运行")
            return
    
    print(f"\n使用RTSP地址: {rtsp_url}")
    print("选择操作:")
    print("1. 连接测试")
    print("2. 实时预览")
    
    choice = input("请选择 (1-2): ").strip()
    
    client = RTSPClient(rtsp_url)
    
    if choice == "1":
        client.test_connection()
    elif choice == "2":
        if not client.test_connection():
            print("连接失败，无法启动预览")
            return
        print()
        client.start_preview()
    else:
        print("无效选择")

if __name__ == "__main__":
    main()
